package com.study.kafka.consumer;

import org.apache.kafka.clients.consumer.*;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.serialization.StringDeserializer;

import java.time.Duration;
import java.util.*;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 消费者
 * KafkaConsumer 是非线程安全的
 *
 * @author: stephen.shen
 * @create: 2019-06-18 11:47
 */
public class DemoKafkaConsumer {

    private final static String brokerList = "localhost:9092";
    private final static String topic = "test-topic";
    private final static String groupId = "group-test";
    private final static AtomicBoolean isRunning = new AtomicBoolean(true);

    private static Properties initConfig() {
        Properties props = new Properties();
        props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList);
        props.put(ConsumerConfig.GROUP_ID_CONFIG, groupId);

        props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG,
                StringDeserializer.class.getName());
        props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG,
                StringDeserializer.class.getName());
        // 禁止默认的自动提交（默认是每5秒提交一次）
        props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, false);
        // 自动提交间隔时间，当ENABLE_AUTO_COMMIT_CONFIG 为false时不生效
        props.put(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG, 100);
        // 每当消费者查找不到所记录的消费位移时(offset不存在或者offset对应的数据已删除)，才会根据这个参数决定从何处开始消费
        // 默认值为"latest"，从分区末尾开始消费
        // "earliest"是从最早开始消费，不一定是从0开始
        // "none"出现差找不到消费位置时抛出NoOffsetForPartitionException
        props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
        // 消费者ID
        props.put(ConsumerConfig.CLIENT_ID_CONFIG, "consumer-1");
        return props;
    }

    public static void main(String[] args) {
        Properties props = initConfig();
        KafkaConsumer<String, String> consumer = new KafkaConsumer<>(props);

        // 订阅主题方式一（不指定分区）
        // 如果多次订阅，以最后一次为准
        //consumer.subscribe(Arrays.asList(topic));
        // 使用正则表达式方式订阅
        // consumer.subscribe(Pattern.compile("test-*"));
        // 添加有再均衡监昕器的订阅
        consumer.subscribe(Arrays.asList(topic), new ConsumerRebalanceListener() {
            /**
             * 在再均衡开始之前和消费者停止读取消息之后被调用
             * @param partitions 再均衡前所分配到的分区
             */
            @Override
            public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
                //commitOffset(currentOffset);
            }

            /**
             * 在重新分配分区之后和消费者开始读取消费之前被调用
             * @param partitions 再均衡后所分配到的分区
             */
            @Override
            public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
//                currentOffset.clear();
//                for (TopicPartition partition : partitions) {
//                    // 定位到最近提交的 offset 位置继续消费
//                    consumer.seek(partition, getOffset(partition));
//                }
            }
        });


        // 订阅主题方式二（指定分区）
        /**
         ArrayList<TopicPartition> partitions = new ArrayList<>();
         // 得到主题的分区信息
         List<PartitionInfo> partitionInfos = consumer.partitionsFor(topic);
         if (partitionInfos != null) {
         for (PartitionInfo partitionInfo : partitionInfos) {
         partitions.add(new TopicPartition(topic, partitionInfo.partition()));
         }
         }
         consumer.assign(partitions);
         **/

        // 总结：subscribe和assign的区别
        // subscribe具备消费者自动再均衡的功能，有ConsumerRebalanceListener类型参数的方法
        // assign不具备消费者自动均衡的功能，没有ConsumerRebalanceListener类型参数的方法

        // 取消订阅
        // 方式一
        //consumer.unsubscribe();
        // 方式二
        //consumer.subscribe(new ArrayList<>());
        // 方式三
        //consumer.assign(new ArrayList<TopicPartition>());

        try {
            while (isRunning.get()) {
                // 超时时间1000毫秒，在消费者的缓冲区里没有可用数据时会发生阻塞，
                // 如果应用线程唯一的工作就是从 Kafka 中拉取并消费消息，
                // 则可以将这个参数设置为最大值 Long.MAX_VALUE
                ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(1000));
                for (TopicPartition tp : records.partitions()) {
                    // 按照分区读取数据
                    List<ConsumerRecord<String, String>> partitionRecords = records.records(tp);
                    for (ConsumerRecord<String, String> record : partitionRecords) {

                        System.out.println("key = " + record.key() + ", value = " + record.value()
                                + ", topic = " + record.topic() + ", partition = " + record.partition()
                                + ", offset = " + record.offset());
                    }
                    long lastConsumedOffset = partitionRecords.get(partitionRecords.size() - 1).offset();
                    // 同步提交指定的消费位移，会阻塞
                    consumer.commitSync(Collections.singletonMap(tp, new OffsetAndMetadata(lastConsumedOffset + 1)));
                    // Last consumed offset is 17
                    System.out.println("Last consumed offset is " + lastConsumedOffset);
                    // Commited offset is 18
                    System.out.println("Commited offset is " + consumer.committed(tp).offset());
                    // Next record offset is 18
                    System.out.println("Next record offset is " + consumer.position(tp));

                    // 异步提交
                    //consumer.commitAsync();
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            try {
                // 消费者正常退出时位移提交
                consumer.commitSync();
            } finally {
                consumer.close();
            }
        }
    }

    /**
     * 指定位移消费
     */
    private void seek(KafkaConsumer consumer) {
        consumer.subscribe(Arrays.asList(topic));
        Set<TopicPartition> assignment = new HashSet<>();
        // 执行seek()方法之前需要先执行一次poll()方法
        // 等到分配到分区之后才可以重置消费位置
        while (assignment.size() == 0) {
            // 参数不能设置为0时，如果设置为0，此方法立刻返回，
            // 那么poll()方法内部进行分区分配的逻辑就会来不及实施
            // 得不到分区信息，assignment会一直是空，会无法继续往下执行
            consumer.poll(Duration.ofMillis(100));
            assignment = consumer.assignment();
        }
        // 获取指定分区的末尾的消息位置
        Map<TopicPartition, Long> offsets = consumer.endOffsets(assignment);
        // 获取指定分区的开头的消息位置
        //Map<TopicPartition,Long> offsets = consumer.beginningOffsets(assignment);
        for (TopicPartition tp : assignment) {
            // 设置指定分区从指定位置开始消费，这里是从末尾开始消费
            consumer.seek(tp, offsets.get(tp));
        }

        // 从分区的开头开始消费, 同beginningOffsets()
        //consumer.seekToBeginning(assignment);
        // 从分区的末尾开始消费, 同endOffsets()
        //consumer.seekToEnd(assignment);

        // 从指定时间开始消费
        /**
         Map<TopicPartition, Long> timestampToSearch = new HashMap<>();
         for (TopicPartition tp : assignment) {
         timestampToSearch.put(tp, System.currentTimeMillis() - 1 * 24 * 3600 * 1000);
         }
         Map<TopicPartition, Long> offsets = consumer.offsetsForTimes(timestampToSearch);
         for (TopicPartition tp : assignment) {
         OffsetAndTimestamp offsetAndTimestamp = offsets.get(tp);
         if (offsetAndTimestamp != null) {
         consumer.seek(tp, offsetAndTimestamp.offset());
         }
         }
         **/

        // 进行消费
        // ......
    }

    //获取某分区的最新 offset
    private static long getOffset(TopicPartition partition) {
        return 0;
    }

    //提交该消费者所有分区的 offset
    private static void commitOffset(Map<TopicPartition, Long> currentOffset) {
    }
}
